import numpy as np

# 1. 生成数据
def generate_vectors(num_vectors=10, dimensions=3):
    """  
    生成向量数据  
    :param num_vectors: 向量数量  
    :param dimensions: 向量维度  
    :return: 随机生成的向量数据  
    """  
    np.random.seed(42)  
    return np.random.rand(num_vectors, dimensions)

data_vectors = generate_vectors(num_vectors=10, dimensions=3)

# 2. LSH 哈希函数实现
class LSH:
    """  
    简单实现基于余弦相似度的LSH  
    """  
    def __init__(self, dimensions, num_hashes):  
        self.dimensions = dimensions  
        self.num_hashes = num_hashes  
        self.hash_planes = np.random.randn(num_hashes, dimensions)  # 随机生成超平面
    
    def hash_function(self, vector):
        """
        对单个向量进行哈希
        :param vector: 输入向量
        :return: 哈希值
        """
        projections = np.dot(self.hash_planes, vector)  # 投影到超平面
        hash_value = ''.join(['1' if p > 0 else '0' for p in projections])  # 大于0为1，否则为0
        return hash_value

    def hash_vectors(self, vectors):
        """
        对所有向量进行哈希
        :param vectors: 输入向量列表
        :return: 哈希表
        """
        hash_table = {}
        for idx, vector in enumerate(vectors):
            hash_value = self.hash_function(vector)
            if hash_value not in hash_table:
                hash_table[hash_value] = []
            hash_table[hash_value].append(idx)
        return hash_table

# 3. 初始化LSH并进行分桶
num_hashes = 4
lsh = LSH(dimensions=3, num_hashes=num_hashes)
hash_table = lsh.hash_vectors(data_vectors)

# 4. 查询相似向量
def query_lsh(query_vector, lsh, hash_table):
    """
    查询与输入向量相似的向量
    :param query_vector: 查询向量
    :param lsh: LSH 实例
    :param hash_table: 哈希表
    :return: 相似向量索引
    """
    query_hash = lsh.hash_function(query_vector)
    return hash_table.get(query_hash, [])

query_vector = np.array([0.5, 0.5, 0.5])  # 查询向量
similar_indices = query_lsh(query_vector, lsh, hash_table)

# 5. 输出结果
print("向量数据:")
print(data_vectors)
print("\n哈希表:")
for bucket, indices in hash_table.items():
    print(f"桶 {bucket}: {indices}")
print("\n查询向量:", query_vector)
print("相似向量索引:", similar_indices)